#include "md5_2.h"
#include <iomanip>
#include <assert.h>
#include <chrono>

using namespace std;
using namespace chrono;

// 重用md5.cpp中的StringProcess函数
Byte *StringProcess(string input, int *n_byte)
{
    // 将输入的字符串转换为Byte为单位的数组
    Byte *blocks = (Byte *)input.c_str();
    int length = input.length();

    // 计算原始消息长度（以比特为单位）
    int bitLength = length * 8;

    // paddingBits: 原始消息需要的padding长度（以bit为单位）
    int paddingBits = bitLength % 512;
    if (paddingBits > 448)
    {
        paddingBits += 512 - (paddingBits - 448);
    }
    else if (paddingBits < 448)
    {
        paddingBits = 448 - paddingBits;
    }
    else if (paddingBits == 448)
    {
        paddingBits = 512;
    }

    // 原始消息需要的padding长度（以Byte为单位）
    int paddingBytes = paddingBits / 8;
    // 创建最终的字节数组
    int paddedLength = length + paddingBytes + 8;
    Byte *paddedMessage = new Byte[paddedLength];

    // 复制原始消息
    memcpy(paddedMessage, blocks, length);

    // 添加填充字节
    paddedMessage[length] = 0x80;
    memset(paddedMessage + length + 1, 0, paddingBytes - 1);

    // 添加消息长度（64比特，小端格式）
    for (int i = 0; i < 8; ++i)
    {
        paddedMessage[length + paddingBytes + i] = ((uint64_t)length * 8 >> (i * 8)) & 0xFF;
    }

    *n_byte = paddedLength;
    return paddedMessage;
}

void SIMD_MD5Hash_2(string input[], bit32 state[2][4])
{
    Byte *paddedMessage[2];
    int *messageLength = new int[2];
    
    // 处理输入字符串
    for (int i = 0; i < 2; i++)
    {
        paddedMessage[i] = StringProcess(input[i], &messageLength[i]);
    }
    
    // 计算每个输入需要的块数
    int n_blocks[2] = {
        messageLength[0] / 64, 
        messageLength[1] / 64
    };
    
    // 找出最大和最小块数
    int max_n_blocks = max(n_blocks[0], n_blocks[1]);
    int min_n_blocks = min(n_blocks[0], n_blocks[1]);

    // 初始化状态
    for (int i = 0; i < 2; i++)
    {
        state[i][0] = 0x67452301;
        state[i][1] = 0xefcdab89;
        state[i][2] = 0x98badcfe;
        state[i][3] = 0x10325476;
    }
    
    // 使用NEON向量初始化状态值 - 使用64位向量(uint32x2_t)
    uint32x2_t a_end = vdup_n_u32(state[0][0]);
    uint32x2_t b_end = vdup_n_u32(state[0][1]);
    uint32x2_t c_end = vdup_n_u32(state[0][2]);
    uint32x2_t d_end = vdup_n_u32(state[0][3]);
    
    uint32x2_t a = a_end;
    uint32x2_t b = b_end;
    uint32x2_t c = c_end; 
    uint32x2_t d = d_end;
    uint32x2_t x[16];

    int over_num = 0;
    
    // 逐块处理
    for (int i = 0; i <= max_n_blocks; i++)
    {
        // 检查是否完成某一输入的处理
        if (i + 1 >= min_n_blocks)
        {
            if (i == n_blocks[0])
            {
                state[0][0] = vget_lane_u32(a_end, 0);
                state[0][1] = vget_lane_u32(b_end, 0);
                state[0][2] = vget_lane_u32(c_end, 0);
                state[0][3] = vget_lane_u32(d_end, 0);
                over_num++;
            }
            if (i == n_blocks[1])
            {
                state[1][0] = vget_lane_u32(a_end, 1);
                state[1][1] = vget_lane_u32(b_end, 1);
                state[1][2] = vget_lane_u32(c_end, 1);
                state[1][3] = vget_lane_u32(d_end, 1);
                over_num++;
            }
            if (over_num == 2)
            {
                break;
            }
        }

        // 为每个16个字处理块
        for (int i1 = 0; i1 < 16; ++i1)
        {
            int offset = 4 * i1 + i * 64;
            // 使用直接加载方式填充2个32位值到NEON寄存器
            x[i1] = (uint32x2_t){
                *reinterpret_cast<uint32_t *>(paddedMessage[0] + offset),
                *reinterpret_cast<uint32_t *>(paddedMessage[1] + offset)
            };
        }

        // MD5算法的四轮循环
        /* Round 1 */
        FF2(a, b, c, d, x[0], s11, 0xd76aa478);
        FF2(d, a, b, c, x[1], s12, 0xe8c7b756);
        FF2(c, d, a, b, x[2], s13, 0x242070db);
        FF2(b, c, d, a, x[3], s14, 0xc1bdceee);
        FF2(a, b, c, d, x[4], s11, 0xf57c0faf);
        FF2(d, a, b, c, x[5], s12, 0x4787c62a);
        FF2(c, d, a, b, x[6], s13, 0xa8304613);
        FF2(b, c, d, a, x[7], s14, 0xfd469501);
        FF2(a, b, c, d, x[8], s11, 0x698098d8);
        FF2(d, a, b, c, x[9], s12, 0x8b44f7af);
        FF2(c, d, a, b, x[10], s13, 0xffff5bb1);
        FF2(b, c, d, a, x[11], s14, 0x895cd7be);
        FF2(a, b, c, d, x[12], s11, 0x6b901122);
        FF2(d, a, b, c, x[13], s12, 0xfd987193);
        FF2(c, d, a, b, x[14], s13, 0xa679438e);
        FF2(b, c, d, a, x[15], s14, 0x49b40821);

        /* Round 2 */
        GG2(a, b, c, d, x[1], s21, 0xf61e2562);
        GG2(d, a, b, c, x[6], s22, 0xc040b340);
        GG2(c, d, a, b, x[11], s23, 0x265e5a51);
        GG2(b, c, d, a, x[0], s24, 0xe9b6c7aa);
        GG2(a, b, c, d, x[5], s21, 0xd62f105d);
        GG2(d, a, b, c, x[10], s22, 0x2441453);
        GG2(c, d, a, b, x[15], s23, 0xd8a1e681);
        GG2(b, c, d, a, x[4], s24, 0xe7d3fbc8);
        GG2(a, b, c, d, x[9], s21, 0x21e1cde6);
        GG2(d, a, b, c, x[14], s22, 0xc33707d6);
        GG2(c, d, a, b, x[3], s23, 0xf4d50d87);
        GG2(b, c, d, a, x[8], s24, 0x455a14ed);
        GG2(a, b, c, d, x[13], s21, 0xa9e3e905);
        GG2(d, a, b, c, x[2], s22, 0xfcefa3f8);
        GG2(c, d, a, b, x[7], s23, 0x676f02d9);
        GG2(b, c, d, a, x[12], s24, 0x8d2a4c8a);

        /* Round 3 */
        HH2(a, b, c, d, x[5], s31, 0xfffa3942);
        HH2(d, a, b, c, x[8], s32, 0x8771f681);
        HH2(c, d, a, b, x[11], s33, 0x6d9d6122);
        HH2(b, c, d, a, x[14], s34, 0xfde5380c);
        HH2(a, b, c, d, x[1], s31, 0xa4beea44);
        HH2(d, a, b, c, x[4], s32, 0x4bdecfa9);
        HH2(c, d, a, b, x[7], s33, 0xf6bb4b60);
        HH2(b, c, d, a, x[10], s34, 0xbebfbc70);
        HH2(a, b, c, d, x[13], s31, 0x289b7ec6);
        HH2(d, a, b, c, x[0], s32, 0xeaa127fa);
        HH2(c, d, a, b, x[3], s33, 0xd4ef3085);
        HH2(b, c, d, a, x[6], s34, 0x4881d05);
        HH2(a, b, c, d, x[9], s31, 0xd9d4d039);
        HH2(d, a, b, c, x[12], s32, 0xe6db99e5);
        HH2(c, d, a, b, x[15], s33, 0x1fa27cf8);
        HH2(b, c, d, a, x[2], s34, 0xc4ac5665);

        /* Round 4 */
        II2(a, b, c, d, x[0], s41, 0xf4292244);
        II2(d, a, b, c, x[7], s42, 0x432aff97);
        II2(c, d, a, b, x[14], s43, 0xab9423a7);
        II2(b, c, d, a, x[5], s44, 0xfc93a039);
        II2(a, b, c, d, x[12], s41, 0x655b59c3);
        II2(d, a, b, c, x[3], s42, 0x8f0ccc92);
        II2(c, d, a, b, x[10], s43, 0xffeff47d);
        II2(b, c, d, a, x[1], s44, 0x85845dd1);
        II2(a, b, c, d, x[8], s41, 0x6fa87e4f);
        II2(d, a, b, c, x[15], s42, 0xfe2ce6e0);
        II2(c, d, a, b, x[6], s43, 0xa3014314);
        II2(b, c, d, a, x[13], s44, 0x4e0811a1);
        II2(a, b, c, d, x[4], s41, 0xf7537e82);
        II2(d, a, b, c, x[11], s42, 0xbd3af235);
        II2(c, d, a, b, x[2], s43, 0x2ad7d2bb);
        II2(b, c, d, a, x[9], s44, 0xeb86d391);

        // 更新状态
        a_end = vadd_u32(a_end, a);
        b_end = vadd_u32(b_end, b);  
        c_end = vadd_u32(c_end, c);
        d_end = vadd_u32(d_end, d);
        
        a = a_end;
        b = b_end;
        c = c_end;
        d = d_end;
    }

    // 处理最终结果，转换字节序
    for (int j = 0; j < 2; j++)
    {
        // 使用单独的32位转换，因为64位向量没有对应的字节逆序操作
        for (int i = 0; i < 4; i++)
        {
            uint32_t value = state[j][i];
            state[j][i] = ((value & 0xff) << 24) |
                          ((value & 0xff00) << 8) |
                          ((value & 0xff0000) >> 8) |
                          ((value & 0xff000000) >> 24);
        }
    }

    // 释放内存
    for (int i = 0; i < 2; i++)
    {
        delete[] paddedMessage[i];
    }
    delete[] messageLength;
}
